// "Copyright [2021] <Copyright Shengkai Lin>"

package cluster.network_devices;

import cluster.optical_devices.*;

module Machine {
    parameters:
        @controlplane_task;
        int numGpus = default(8);
        int numNvlink_port = default(6);
        int my_index;
        int nic_type = default(0);  // udp-roce = 0; ib = 1
        bool connectToOXC = false;
        bool need_wait_topo_machine = default(true);

    gates:
        inout out_port[numGpus];  // Data plane
        inout controller_port;  // Control plane

    submodules:
        cpu[numGpus+1] : Cpu {
            parameters:
                my_index = index;
                my_global_index = index + my_index*numGpus;
                need_wait_topo = need_wait_topo_machine;
        }
        gpu[numGpus] : Gpu {
            parameters:
                numNvlink_port = numNvlink_port;
                my_global_index = index + my_index*numGpus;
                my_index = index;
        }
        nic[numGpus] : <default("SimpleNic")> like INic{
            parameters:
                my_index = index;
                my_global_index = index + my_index*numGpus;
        }

        opticaltransfer[numGpus] : OpticalTransceiverBuf if connectToOXC{
            parameters:
                **.my_global_index = index + my_index*numGpus;
                **.my_index = index;
        }

    connections allowunconnected:
        // Only connect cpu0 to the controller directly.
        cpu[0].controller_port$o --> controller_port$o;
        cpu[0].controller_port$i <-- controller_port$i;

        for i=1..(numGpus) {
            cpu[i].controller_port <-->Pcie<--> cpu[0].cpu_port++;
            gpu[i-1].manager_port <-->Pcie<--> cpu[i].gpu_port;
            gpu[i-1].socket_port <-->Pcie<--> nic[i-1].gpu_port;
            nic[i-1].out_port <-->FiberNoDealy<--> opticaltransfer[i-1].nic_port if connectToOXC == true;
            opticaltransfer[i-1].out_port$o --> out_port$o++ if connectToOXC == true;
            opticaltransfer[i-1].out_port$i <-- out_port$i++ if connectToOXC == true;
            nic[i-1].out_port <--> out_port++ if connectToOXC == false;
        }

        // The connections of NVlink
        // Connect the basic pattern of nvlink
        // only 1 connection between 2 peers
        gpu[0].nvlink_port[1] <-->Nvlink<--> gpu[1].nvlink_port[3] if numNvlink_port>=4;
        gpu[1].nvlink_port[1] <-->Nvlink<--> gpu[2].nvlink_port[3] if numNvlink_port>=4;
        gpu[2].nvlink_port[1] <-->Nvlink<--> gpu[3].nvlink_port[3] if numNvlink_port>=4;
        gpu[3].nvlink_port[1] <-->Nvlink<--> gpu[0].nvlink_port[3] if numNvlink_port>=4;

        gpu[0].nvlink_port[2] <-->Nvlink<--> gpu[2].nvlink_port[2] if numNvlink_port>=4;
        gpu[1].nvlink_port[2] <-->Nvlink<--> gpu[3].nvlink_port[2] if numNvlink_port>=4;

        gpu[4].nvlink_port[1] <-->Nvlink<--> gpu[5].nvlink_port[3] if numNvlink_port>=4;
        gpu[5].nvlink_port[1] <-->Nvlink<--> gpu[6].nvlink_port[3] if numNvlink_port>=4;
        gpu[6].nvlink_port[1] <-->Nvlink<--> gpu[7].nvlink_port[3] if numNvlink_port>=4;
        gpu[7].nvlink_port[1] <-->Nvlink<--> gpu[4].nvlink_port[3] if numNvlink_port>=4;

        gpu[4].nvlink_port[2] <-->Nvlink<--> gpu[6].nvlink_port[2] if numNvlink_port>=4;
        gpu[5].nvlink_port[2] <-->Nvlink<--> gpu[7].nvlink_port[2] if numNvlink_port>=4;

        for i=0 .. 3 {
            gpu[i].nvlink_port[0] <-->Nvlink<--> gpu[i+4].nvlink_port[0] if numNvlink_port>=4;
        }

        // Connect the extra nvlink port for ring.
        gpu[0].nvlink_port[4] <-->Nvlink<--> gpu[4].nvlink_port[5] if numNvlink_port == 6;
        gpu[4].nvlink_port[4] <-->Nvlink<--> gpu[7].nvlink_port[5] if numNvlink_port == 6;
        gpu[7].nvlink_port[4] <-->Nvlink<--> gpu[6].nvlink_port[5] if numNvlink_port == 6;
        gpu[6].nvlink_port[4] <-->Nvlink<--> gpu[5].nvlink_port[5] if numNvlink_port == 6;
        gpu[5].nvlink_port[4] <-->Nvlink<--> gpu[1].nvlink_port[5] if numNvlink_port == 6;
        gpu[1].nvlink_port[4] <-->Nvlink<--> gpu[2].nvlink_port[5] if numNvlink_port == 6;
        gpu[2].nvlink_port[4] <-->Nvlink<--> gpu[3].nvlink_port[5] if numNvlink_port == 6;
        gpu[3].nvlink_port[4] <-->Nvlink<--> gpu[0].nvlink_port[5] if numNvlink_port == 6;

        // Connect peci port.
        gpu[0].pcie_port++ <-->Pcie_real<--> gpu[5].pcie_port++;
        gpu[0].pcie_port++ <-->Pcie_real<--> gpu[6].pcie_port++;
        gpu[0].pcie_port++ <-->Pcie_real<--> gpu[7].pcie_port++;
        gpu[1].pcie_port++ <-->Pcie_real<--> gpu[4].pcie_port++;
        gpu[1].pcie_port++ <-->Pcie_real<--> gpu[6].pcie_port++;
        gpu[1].pcie_port++ <-->Pcie_real<--> gpu[7].pcie_port++;
        gpu[2].pcie_port++ <-->Pcie_real<--> gpu[4].pcie_port++;
        gpu[2].pcie_port++ <-->Pcie_real<--> gpu[5].pcie_port++;
        //gpu[2].pcie_port++ <-->Pcie_real<--> gpu[6].pcie_port++;
        gpu[2].pcie_port++ <-->Pcie_real<--> gpu[7].pcie_port++;
        gpu[3].pcie_port++ <-->Pcie_real<--> gpu[4].pcie_port++;
        gpu[3].pcie_port++ <-->Pcie_real<--> gpu[5].pcie_port++;
        gpu[3].pcie_port++ <-->Pcie_real<--> gpu[6].pcie_port++;
        //gpu[3].pcie_port++ <-->Pcie_real<--> gpu[7].pcie_port++;
}